{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ccfefb2a",
   "metadata": {},
   "source": [
    "给高loss的数据加权,提高这些数据被采样中的概率.\n",
    "\n",
    "为了缓解过拟合,降低这些数据的lr,这里以削减loss的方式实现."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "91162df8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAT8klEQVR4nO3dfWxT570H8K+dxIaQHGeBxl4u8QWpbCyXl7YBktP+san1yGhUjTa72irUZQhRlTmoNBPSIrVBbTelYtLadaPhnw1aTR1T0GVTI/oShTZowklK2lyFULJWoksuYHstynFeiJ34/O4fKGc1BBaT4MfG3490JPw8j49/5wF/OS8+tk1EBERECthVF0BE2YsBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyigLoAMHDmDFihVYtGgRKisr0dPTo6oUIlJESQD9+c9/RkNDA/bt24ePPvoI69evR3V1NcLhsIpyiEgRm4qbUSsrK7Fx40b87ne/AwCYpomysjLs3r0bP//5z1NdDhEpkpvqF4zFYujt7UVjY6PVZrfb4fP5EAgEZn1ONBpFNBq1HpumicuXL2Pp0qWw2Wy3vWYiSo6IYHR0FKWlpbDbb3yglfIA+uKLLxCPx+F2uxPa3W43zp07N+tzmpub8fzzz6eiPCJaQMPDw1i+fPkN+1MeQLeisbERDQ0N1mPDMOD1ejE8PAxN0xRWRkSziUQiKCsrQ2Fh4U3HpTyAli1bhpycHIRCoYT2UCgEj8cz63OcTiecTud17ZqmMYCI0ti/O0WS8qtgDocDFRUV6OjosNpM00RHRwd0XU91OUSkkJJDsIaGBtTV1WHDhg3YtGkTXnnlFYyPj2P79u0qyiEiRZQE0A9/+EP885//RFNTE4LBIO655x688847152YJqI7m5LPAc1XJBKBy+WCYRg8B0SUhub6HuW9YESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlIm6QA6efIkHnnkEZSWlsJms+Evf/lLQr+IoKmpCV//+texePFi+Hw+fPrppwljLl++jG3btkHTNBQVFWHHjh0YGxub14YQUeZJOoDGx8exfv16HDhwYNb+/fv349VXX8XBgwfR3d2NJUuWoLq6GpOTk9aYbdu2YWBgAO3t7Whra8PJkyfx5JNP3vpWEFFmknkAIMeOHbMem6YpHo9HfvWrX1ltIyMj4nQ65U9/+pOIiJw9e1YAyIcffmiNefvtt8Vms8mFCxfm9LqGYQgAMQxjPuUT0W0y1/fogp4DOn/+PILBIHw+n9XmcrlQWVmJQCAAAAgEAigqKsKGDRusMT6fD3a7Hd3d3bOuNxqNIhKJJCxElPkWNICCwSAAwO12J7S73W6rLxgMoqSkJKE/NzcXxcXF1phrNTc3w+VyWUtZWdlClk1EimTEVbDGxkYYhmEtw8PDqksiogWwoAHk8XgAAKFQKKE9FApZfR6PB+FwOKF/enoaly9ftsZcy+l0QtO0hIWIMt+CBtDKlSvh8XjQ0dFhtUUiEXR3d0PXdQCArusYGRlBb2+vNebEiRMwTROVlZULWQ4RpbncZJ8wNjaGzz77zHp8/vx59PX1obi4GF6vF3v27MEvfvELrFq1CitXrsRzzz2H0tJSbN26FQDwrW99C9/73vewc+dOHDx4EFNTU6ivr8ePfvQjlJaWLtiGEVEGSPby2vvvvy8Arlvq6upE5Oql+Oeee07cbrc4nU556KGHZHBwMGEdX375pTz++ONSUFAgmqbJ9u3bZXR0dMEv8RGRGnN9j9pERBTm3y2JRCJwuVwwDIPng4jS0FzfoxlxFYyI7kwMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEiZpH+Wh2ihiJgwhgcQj05Ybfl3/ScWudyw2WwKK6NUYQCRMmKa+L/u/8GVyxestrKq/8aidW6FVVEq8RCM0oppTqsugVKIAURpRcy46hIohRhAlFYkzj2gbMIAorTCPaDswgCitCI8B5RVGECUVsx4HICoLoNSJKkAam5uxsaNG1FYWIiSkhJs3boVg4ODCWMmJyfh9/uxdOlSFBQUoLa2FqFQKGHM0NAQampqkJ+fj5KSEuzduxfT0/yfL9vYbDbYc50JbfHYBPMniyQVQJ2dnfD7/ejq6kJ7ezumpqawefNmjI+PW2OeeeYZvPXWW2htbUVnZycuXryIxx57zOqPx+OoqalBLBbDqVOn8Prrr+Pw4cNoampauK2izGCzYVGRJ6FpciQIJlAWkXkIh8MCQDo7O0VEZGRkRPLy8qS1tdUa88knnwgACQQCIiJy/PhxsdvtEgwGrTEtLS2iaZpEo9E5va5hGAJADMOYT/mkmGnG5XznG9JzcKe1nDn6opjxadWl0TzN9T06r3NAhmEAAIqLiwEAvb29mJqags/ns8asXr0aXq8XgUAAABAIBLB27Vq43f/6tGt1dTUikQgGBgZmfZ1oNIpIJJKw0J3BZueH8bPZLQeQaZrYs2cPHnjgAaxZswYAEAwG4XA4UFRUlDDW7XYjGAxaY74aPjP9M32zaW5uhsvlspaysrJbLZvSig22nBzVRZBCtxxAfr8fZ86cwZEjRxaynlk1NjbCMAxrGR4evu2vSanBPaDsdkt/+/X19Whra8PJkyexfPlyq93j8SAWi2FkZCRhLygUCsHj8Vhjenp6EtY3c5VsZsy1nE4nnE7nrH2U2ewMoKyW1B6QiKC+vh7Hjh3DiRMnsHLlyoT+iooK5OXloaOjw2obHBzE0NAQdF0HAOi6jv7+foTDYWtMe3s7NE1DeXn5fLaFMhAPwbJbUv/9+P1+vPnmm/jrX/+KwsJC65yNy+XC4sWL4XK5sGPHDjQ0NKC4uBiapmH37t3QdR1VVVUAgM2bN6O8vBxPPPEE9u/fj2AwiGeffRZ+v597OVnIZmcAZbOkAqilpQUA8J3vfCeh/dChQ/jJT34CAHj55Zdht9tRW1uLaDSK6upqvPbaa9bYnJwctLW1YdeuXdB1HUuWLEFdXR1eeOGF+W0JZSSb7ZqdcBGICPh1ZNnBJiIZ96mvSCQCl8sFwzCgaZrqcmgegv/7Hoa7jlqPFxf/B8ofbYQ916GwKpqvub5HeS8YpRURE2KaqsugFGEAUXoREyIMoGzBAKK0IqYADKCswQCitCLcA8oqDCBKLzwHlFUYQKTWNZfhxTT5taxZhAFESi1yuRPuB5ueHMXUFX7bQbZgAJFS9tw84Cu/gipi8iR0FmEAkVK8Gz67MYBIKVtODn8HPosxgEgp3oya3RhApNTVAOIeULZiAJFSNnsO8yeLMYBIqavfiMgEylYMIFLr2hPQcvWbNyk7MIAo7Uicn4TOFgwgSjti8me6swUDiNKOyQDKGgwgSjs8BMseDCBKOzwEyx4MIEozwj2gLMIAorTDc0DZgwFEStnznMjLdyW0RY3wDUbTnYYBRErZ7bnIdSxOaIvHJhRVQ6nGL2Oh2y4ajeLKlSuz9pnTMcRNuWZ8DCMjIzdcX35+PhwO/nDhnYABRLfdG2+8geeff37WvrxcO5oe34j/+s+lVtvRo61o2b7vhut75ZVX8IMf/GDB66TUYwDRbTc2NoYLFy7M2peXa8foxL347Mp9uDxVirvy/oHRsb/fcDwATEzwEO1OkdQ5oJaWFqxbtw6apkHTNOi6jrffftvqn5ychN/vx9KlS1FQUIDa2lqEQqGEdQwNDaGmpgb5+fkoKSnB3r17MT3Nqx7ZyjQFfx9bh88mrgbQ3ycq8Y/JctVlUYokFUDLly/HSy+9hN7eXpw+fRoPPvggvv/972NgYAAA8Mwzz+Ctt95Ca2srOjs7cfHiRTz22GPW8+PxOGpqahCLxXDq1Cm8/vrrOHz4MJqamhZ2qyhjiADGlAsz/xQFdozHi5TWRKmTVAA98sgjePjhh7Fq1Sp84xvfwC9/+UsUFBSgq6sLhmHg97//PX7961/jwQcfREVFBQ4dOoRTp06hq6sLAPDee+/h7Nmz+OMf/4h77rkHW7ZswYsvvogDBw4gFovdlg2k9GaKoCT3U+TaogAEDtsVeBznVZdFKXLL54Di8ThaW1sxPj4OXdfR29uLqakp+Hw+a8zq1avh9XoRCARQVVWFQCCAtWvXwu12W2Oqq6uxa9cuDAwM4N57702qhnPnzqGgoOBWN4FS5NrD8Gv1fvweXP/4HCPTJSjOCyJ46dxNx1+8eBFnz55dyBJpgY2Njc1pXNIB1N/fD13XMTk5iYKCAhw7dgzl5eXo6+uDw+FAUVFRwni3241gMAgACAaDCeEz0z/TdyPRaBTRaNR6HIlc/eE6wzB4/igD3OgS/IzOvs8BfD7n9U1MTNz0Mj2pNz4+PqdxSQfQN7/5TfT19cEwDBw9ehR1dXXo7OxMusBkNDc3z3oZt7KyEpqm3dbXpvnr7u5e0PXdfffduP/++xd0nbSwZnYS/p2kPwntcDhw9913o6KiAs3NzVi/fj1+85vfwOPxIBa7/gNkoVAIHo8HAODxeK7bHZ95PDNmNo2NjTAMw1qGh4eTLZuI0tC8b8UwTRPRaBQVFRXIy8tDR0eH1Tc4OIihoSHoug4A0HUd/f39CIf/da9Pe3s7NE1DefmNL706nU7r0v/MQkSZL6lDsMbGRmzZsgVerxejo6N488038cEHH+Ddd9+Fy+XCjh070NDQgOLiYmiaht27d0PXdVRVVQEANm/ejPLycjzxxBPYv38/gsEgnn32Wfj9fjidztuygUSUvpIKoHA4jB//+Me4dOkSXC4X1q1bh3fffRff/e53AQAvv/wy7HY7amtrEY1GUV1djddee816fk5ODtra2rBr1y7ouo4lS5agrq4OL7zwwsJuFaWVmT3YhcL7wO4cNsnA30CJRCJwuVwwDIOHYxlgdHR0Qa9aFRcXY8mSJQu2Plp4c32P8l4wuu0KCwtRWFiougxKQ/w+ICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKRMruoCboWIAAAikYjiSohoNjPvzZn36o1kZAB9+eWXAICysjLFlRDRzYyOjsLlct2wPyMDqLi4GAAwNDR0042jRJFIBGVlZRgeHoamaarLyQics1sjIhgdHUVpaelNx2VkANntV09duVwu/qO4BZqmcd6SxDlL3lx2DngSmoiUYQARkTIZGUBOpxP79u2D0+lUXUpG4bwlj3N2e9nk310nIyK6TTJyD4iI7gwMICJShgFERMowgIhImYwMoAMHDmDFihVYtGgRKisr0dPTo7okZZqbm7Fx40YUFhaipKQEW7duxeDgYMKYyclJ+P1+LF26FAUFBaitrUUoFEoYMzQ0hJqaGuTn56OkpAR79+7F9PR0KjdFmZdeegk2mw179uyx2jhnKSIZ5siRI+JwOOQPf/iDDAwMyM6dO6WoqEhCoZDq0pSorq6WQ4cOyZkzZ6Svr08efvhh8Xq9MjY2Zo156qmnpKysTDo6OuT06dNSVVUl999/v9U/PT0ta9asEZ/PJx9//LEcP35cli1bJo2NjSo2KaV6enpkxYoVsm7dOnn66aetds5ZamRcAG3atEn8fr/1OB6PS2lpqTQ3NyusKn2Ew2EBIJ2dnSIiMjIyInl5edLa2mqN+eSTTwSABAIBERE5fvy42O12CQaD1piWlhbRNE2i0WhqNyCFRkdHZdWqVdLe3i7f/va3rQDinKVORh2CxWIx9Pb2wufzWW12ux0+nw+BQEBhZenDMAwA/7pht7e3F1NTUwlztnr1ani9XmvOAoEA1q5dC7fbbY2prq5GJBLBwMBACqtPLb/fj5qamoS5AThnqZRRN6N+8cUXiMfjCX/pAOB2u3Hu3DlFVaUP0zSxZ88ePPDAA1izZg0AIBgMwuFwoKioKGGs2+1GMBi0xsw2pzN9d6IjR47go48+wocffnhdH+csdTIqgOjm/H4/zpw5g7/97W+qS0lrw8PDePrpp9He3o5FixapLierZdQh2LJly5CTk3Pd1YhQKASPx6OoqvRQX1+PtrY2vP/++1i+fLnV7vF4EIvFMDIykjD+q3Pm8XhmndOZvjtNb28vwuEw7rvvPuTm5iI3NxednZ149dVXkZubC7fbzTlLkYwKIIfDgYqKCnR0dFhtpmmio6MDuq4rrEwdEUF9fT2OHTuGEydOYOXKlQn9FRUVyMvLS5izwcFBDA0NWXOm6zr6+/sRDoetMe3t7dA0DeXl5anZkBR66KGH0N/fj76+PmvZsGEDtm3bZv2Zc5Yiqs+CJ+vIkSPidDrl8OHDcvbsWXnyySelqKgo4WpENtm1a5e4XC754IMP5NKlS9YyMTFhjXnqqafE6/XKiRMn5PTp06Lruui6bvXPXFLevHmz9PX1yTvvvCN33XVXVl1S/upVMBHOWapkXACJiPz2t78Vr9crDodDNm3aJF1dXapLUgbArMuhQ4esMVeuXJGf/vSn8rWvfU3y8/Pl0UcflUuXLiWs5/PPP5ctW7bI4sWLZdmyZfKzn/1MpqamUrw16lwbQJyz1ODXcRCRMhl1DoiI7iwMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISJn/Bwc8rOY9ZOyNAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import gym\n",
    "\n",
    "\n",
    "#定义环境\n",
    "class MyWrapper(gym.Wrapper):\n",
    "\n",
    "    def __init__(self):\n",
    "        env = gym.make('CartPole-v1', render_mode='rgb_array')\n",
    "        super().__init__(env)\n",
    "        self.env = env\n",
    "        self.step_n = 0\n",
    "\n",
    "    def reset(self):\n",
    "        state, _ = self.env.reset()\n",
    "        self.step_n = 0\n",
    "        return state\n",
    "\n",
    "    def step(self, action):\n",
    "        state, reward, terminated, truncated, info = self.env.step(action)\n",
    "        over = terminated or truncated\n",
    "\n",
    "        #限制最大步数\n",
    "        self.step_n += 1\n",
    "        if self.step_n >= 200:\n",
    "            over = True\n",
    "        \n",
    "        #没坚持到最后,扣分\n",
    "        if over and self.step_n < 200:\n",
    "            reward = -1000\n",
    "\n",
    "        return state, reward, over\n",
    "\n",
    "    #打印游戏图像\n",
    "    def show(self):\n",
    "        from matplotlib import pyplot as plt\n",
    "        plt.figure(figsize=(3, 3))\n",
    "        plt.imshow(self.env.render())\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "env = MyWrapper()\n",
    "\n",
    "env.reset()\n",
    "\n",
    "env.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ecfbe912",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Sequential(\n",
       "   (0): Linear(in_features=4, out_features=64, bias=True)\n",
       "   (1): ReLU()\n",
       "   (2): Linear(in_features=64, out_features=64, bias=True)\n",
       "   (3): ReLU()\n",
       "   (4): Linear(in_features=64, out_features=2, bias=True)\n",
       " ),\n",
       " Sequential(\n",
       "   (0): Linear(in_features=4, out_features=64, bias=True)\n",
       "   (1): ReLU()\n",
       "   (2): Linear(in_features=64, out_features=64, bias=True)\n",
       "   (3): ReLU()\n",
       "   (4): Linear(in_features=64, out_features=2, bias=True)\n",
       " ))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "#定义模型,评估状态下每个动作的价值\n",
    "model = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 2),\n",
    ")\n",
    "\n",
    "#延迟更新的模型,用于计算target\n",
    "model_delay = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 2),\n",
    ")\n",
    "\n",
    "#复制参数\n",
    "model_delay.load_state_dict(model.state_dict())\n",
    "\n",
    "model, model_delay"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "84cbf0ff",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "d:\\appDir\\python3.10\\lib\\site-packages\\gym\\utils\\passive_env_checker.py:233: DeprecationWarning: `np.bool8` is a deprecated alias for `np.bool_`.  (Deprecated NumPy 1.24)\n",
      "  if not isinstance(terminated, (bool, np.bool8)):\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "-991.0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython import display\n",
    "import random\n",
    "\n",
    "\n",
    "#玩一局游戏并记录数据\n",
    "def play(show=False):\n",
    "    data = []\n",
    "    reward_sum = 0\n",
    "\n",
    "    state = env.reset()\n",
    "    over = False\n",
    "    while not over:\n",
    "        action = model(torch.FloatTensor(state).reshape(1, 4)).argmax().item()\n",
    "        if random.random() < 0.1:\n",
    "            action = env.action_space.sample()\n",
    "\n",
    "        next_state, reward, over = env.step(action)\n",
    "\n",
    "        data.append((state, action, reward, next_state, over))\n",
    "        reward_sum += reward\n",
    "\n",
    "        state = next_state\n",
    "\n",
    "        if show:\n",
    "            display.clear_output(wait=True)\n",
    "            env.show()\n",
    "\n",
    "    return data, reward_sum\n",
    "\n",
    "\n",
    "play()[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c79f77f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Administrator\\AppData\\Local\\Temp\\ipykernel_17216\\919455838.py:35: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at ..\\torch\\csrc\\utils\\tensor_new.cpp:248.)\n",
      "  state = torch.FloatTensor([i[0] for i in data]).reshape(-1, 4)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(201,\n",
       " (array([-0.03484228,  0.02932921,  0.03570003,  0.04746766], dtype=float32),\n",
       "  1,\n",
       "  1.0,\n",
       "  array([-0.03425569,  0.22392155,  0.03664938, -0.23374128], dtype=float32),\n",
       "  False))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#数据池\n",
    "class Pool:\n",
    "\n",
    "    def __init__(self):\n",
    "        self.pool = []\n",
    "        self.prob = []\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.pool)\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return self.pool[i]\n",
    "\n",
    "    #更新动作池\n",
    "    def update(self):\n",
    "        #每次更新不少于N条新数据\n",
    "        old_len = len(self.pool)\n",
    "        while len(pool) - old_len < 200:\n",
    "            data = play()[0]\n",
    "            self.pool.extend(data)\n",
    "            #维护概率表\n",
    "            self.prob.extend([1.0] * len(data))\n",
    "\n",
    "        #只保留最新的N条数据\n",
    "        self.pool = self.pool[-2_0000:]\n",
    "        self.prob = self.prob[-2_0000:]\n",
    "\n",
    "    #获取一批数据样本\n",
    "    def sample(self):\n",
    "        idx = torch.FloatTensor(self.prob).clamp(0.1, 1.0).multinomial(\n",
    "            num_samples=64, replacement=False)\n",
    "\n",
    "        data = [self.pool[i] for i in idx]\n",
    "\n",
    "        state = torch.FloatTensor([i[0] for i in data]).reshape(-1, 4)\n",
    "        action = torch.LongTensor([i[1] for i in data]).reshape(-1, 1)\n",
    "        reward = torch.FloatTensor([i[2] for i in data]).reshape(-1, 1)\n",
    "        next_state = torch.FloatTensor([i[3] for i in data]).reshape(-1, 4)\n",
    "        over = torch.LongTensor([i[4] for i in data]).reshape(-1, 1)\n",
    "\n",
    "        return idx, state, action, reward, next_state, over\n",
    "\n",
    "\n",
    "pool = Pool()\n",
    "pool.update()\n",
    "pool.sample()\n",
    "\n",
    "len(pool), pool[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "018c4b1e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 407 [0.7013702392578125] -958.1\n"
     ]
    }
   ],
   "source": [
    "#训练\n",
    "def train():\n",
    "    model.train()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=2e-4)\n",
    "    loss_fn = torch.nn.MSELoss(reduction='none')\n",
    "\n",
    "    #共更新N轮数据\n",
    "    for epoch in range(100):\n",
    "        pool.update()\n",
    "\n",
    "        #每次更新数据后,训练N次\n",
    "        for i in range(200):\n",
    "\n",
    "            #采样N条数据\n",
    "            idx, state, action, reward, next_state, over = pool.sample()\n",
    "\n",
    "            #计算value\n",
    "            value = model(state).gather(dim=1, index=action)\n",
    "\n",
    "            #计算target\n",
    "            with torch.no_grad():\n",
    "                target = model_delay(next_state)\n",
    "            target = target.max(dim=1)[0].reshape(-1, 1)\n",
    "            target = target * 0.99 * (1 - over) + reward\n",
    "\n",
    "            #根据概率缩放loss\n",
    "            r = torch.FloatTensor([pool.prob[i] for i in idx])\n",
    "            r = (1 - r).clamp(0.1, 1.0).reshape(-1, 1)\n",
    "\n",
    "            loss = loss_fn(value, target)\n",
    "            (loss * r).mean(0).backward()\n",
    "            optimizer.step()\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            #根据loss调整数据权重\n",
    "            for i, j in zip(idx.tolist(),\n",
    "                            loss.abs().sigmoid().flatten().tolist()):\n",
    "                pool.prob[i] = j\n",
    "\n",
    "        #复制参数\n",
    "        if (epoch + 1) % 5 == 0:\n",
    "            model_delay.load_state_dict(model.state_dict())\n",
    "\n",
    "        if epoch % 100 == 0:\n",
    "            test_result = sum([play()[-1] for _ in range(20)]) / 20\n",
    "            print(epoch, len(pool), pool.prob[::5000], test_result)\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a101d0e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAATtUlEQVR4nO3df0xb97kG8McGbELgmEGKXS74JletmnJJ0o0kcBpdbWrdsA5Vy4qutinqWBUlamqiJkzRhtQma7aJiv3RrVtK/rhdUmnKmJjEqnDTdgxaoqkmNHRIhCSovWsu3Da2lzKOgQYb8Hv/aDirE6A4YH9x8nykE8Xn+/r4PQf74fzggEVEBEREClhVN0BEdy4GEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKaMsgI4ePYq1a9ciMzMT5eXl6OnpUdUKESmiJIB+//vfo66uDocPH8Z7772HTZs2obKyEsFgUEU7RKSIRcXNqOXl5diyZQt+/etfAwCi0SiKi4uxb98+/OhHP0p2O0SkSHqyXzASiaC3txf19fXmPKvVCo/HA5/PN+dzwuEwwuGw+TgajWJkZAT5+fmwWCwJ75mI4iMiGBsbQ2FhIazW+Q+0kh5AV69exczMDJxOZ8x8p9OJS5cuzfmchoYGPP/888loj4iW0fDwMIqKiuYdT3oA3Yr6+nrU1dWZjw3DgNvtxvDwMDRNU9gZEc0lFAqhuLgYOTk5C9YlPYDWrFmDtLQ0BAKBmPmBQAAul2vO59jtdtjt9pvma5rGACJawb7oFEnSr4LZbDaUlZWho6PDnBeNRtHR0QFd15PdDhEppOQQrK6uDjU1Ndi8eTO2bt2KX/ziF5iYmMCTTz6poh0iUkRJAH3729/G3//+dxw6dAh+vx8PPPAA3njjjZtOTBPR7U3JzwEtVSgUgsPhgGEYPAdEtAIt9jPKe8GISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkTNwBdObMGTz22GMoLCyExWLBH//4x5hxEcGhQ4dw9913Y9WqVfB4PHj//fdjakZGRrBz505omobc3Fzs2rUL4+PjS1oRIko9cQfQxMQENm3ahKNHj8453tjYiJdeegnHjh3D2bNnsXr1alRWVmJyctKs2blzJwYGBtDe3o62tjacOXMGe/bsufW1IKLUJEsAQFpbW83H0WhUXC6X/PznPzfnjY6Oit1ul9/97nciInLhwgUBIO+++65Z8/rrr4vFYpGPPvpoUa9rGIYAEMMwltI+ESXIYj+jy3oO6MMPP4Tf74fH4zHnORwOlJeXw+fzAQB8Ph9yc3OxefNms8bj8cBqteLs2bNzLjccDiMUCsVMRJT6ljWA/H4/AMDpdMbMdzqd5pjf70dBQUHMeHp6OvLy8syaGzU0NMDhcJhTcXHxcrZNRIqkxFWw+vp6GIZhTsPDw6pbIqJlsKwB5HK5AACBQCBmfiAQMMdcLheCwWDM+PT0NEZGRsyaG9ntdmiaFjMRUepb1gBat24dXC4XOjo6zHmhUAhnz56FrusAAF3XMTo6it7eXrOms7MT0WgU5eXly9kOEa1w6fE+YXx8HB988IH5+MMPP0RfXx/y8vLgdruxf/9+/PSnP8W9996LdevW4bnnnkNhYSF27NgBALj//vvx9a9/Hbt378axY8cwNTWF2tpafOc730FhYeGyrRgRpYB4L6+99dZbAuCmqaamRkQ+uxT/3HPPidPpFLvdLg8//LAMDg7GLOOTTz6R7373u5KdnS2apsmTTz4pY2Njy36Jj4jUWOxn1CIiojD/bkkoFILD4YBhGDwfRLQCLfYzmhJXwYjo9sQAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJSJ+8/yEMVj6toYQv93EZ/98RTAYk2Dw70BaRl2tY3RisAAooQKG0H8rfMVzAZQWkYm/v0/DzOACAAPwSjJ5HP/EjGAKOlS8E/RUYIwgCj5GEB0HQOIFGAA0WcYQJRkwkMwMjGAKPkYQHRdXAHU0NCALVu2ICcnBwUFBdixYwcGBwdjaiYnJ+H1epGfn4/s7GxUV1cjEAjE1AwNDaGqqgpZWVkoKCjAwYMHMT09vfS1odTAAKLr4gqgrq4ueL1edHd3o729HVNTU9i+fTsmJibMmgMHDuDUqVNoaWlBV1cXPv74Yzz++OPm+MzMDKqqqhCJRPDOO+/g1VdfxYkTJ3Do0KHlWytauQQQngOiWbIEwWBQAEhXV5eIiIyOjkpGRoa0tLSYNRcvXhQA4vP5RETk9OnTYrVaxe/3mzVNTU2iaZqEw+FFva5hGAJADMNYSvuUBGNXPpCeY3uk59hu6Tm2W879V61MXB1W3RYl2GI/o0s6B2QYBgAgLy8PANDb24upqSl4PB6zZv369XC73fD5fAAAn8+HDRs2wOl0mjWVlZUIhUIYGBiY83XC4TBCoVDMRCmMh2B03S0HUDQaxf79+7Ft2zaUlpYCAPx+P2w2G3Jzc2NqnU4n/H6/WfP58Jkdnx2bS0NDAxwOhzkVFxffatuknPAQjEy3HEBerxfnz59Hc3PzcvYzp/r6ehiGYU7Dw8MJf01KIO4B0XW3dDNqbW0t2tracObMGRQVFZnzXS4XIpEIRkdHY/aCAoEAXC6XWdPT0xOzvNmrZLM1N7Lb7bDbefPibYMBRNfFtQckIqitrUVrays6Ozuxbt26mPGysjJkZGSgo6PDnDc4OIihoSHoug4A0HUd/f39CAaDZk17ezs0TUNJSclS1oVSBA/BaFZce0BerxcnT57Ea6+9hpycHPOcjcPhwKpVq+BwOLBr1y7U1dUhLy8PmqZh37590HUdFRUVAIDt27ejpKQETzzxBBobG+H3+/Hss8/C6/VyL+dOwT0gui6uAGpqagIAfO1rX4uZf/z4cXz/+98HALz44ouwWq2orq5GOBxGZWUlXn75ZbM2LS0NbW1t2Lt3L3Rdx+rVq1FTU4MjR44sbU0oJYjwVgz6J4uk4LshFArB4XDAMAxomqa6HVrAuP9/cPG1Rpi/ETEtA/dVHUDO3feobYwSarGfUd4LRgqk3Pc8ShAGECVdCu50U4IwgCj5GEB0HQOIkkzAQzCaxQCipBOJqm6BVggGECWWxfLZNEsAiTKA6DMMIEqojFU5SM/MNh9LdBrhUHCBZ9CdhAFEiWVNg8US+zaT6IyiZmilYQBRQllgASxfXEd3JgYQJZbF/IfoJgwgSjAL44fmxQCixLMwgmhuDCBKLIsFPASj+TCAKKF4EpoWwgCixOIeEC2AAUQJx/ih+TCAKLFuvBWD6HMYQJRQjB5aCAOIEot7QLQABhAlGE9C0/wYQJRwjB+aDwOIEsrCQzBaAAOIEozhQ/NjAFFi8W54WkBcfxmVaC6ffvopIpHInGPR6QiiN/wK1mvXrmF0dHTe5WVnZyM9nW/NOwG/yrRkR44cwW9/+9s5x2zpVjTu2gb3XTnmvMbGRpw6+/Sc9VarFa2trSgrK0tIr7SyMIBoyUZHR/HRRx/NOWZLT8NE2IqLEzrGZvJQZB/EqNE9b73Vap13b4puP3GdA2pqasLGjRuhaRo0TYOu63j99dfN8cnJSXi9XuTn5yM7OxvV1dUIBAIxyxgaGkJVVRWysrJQUFCAgwcPYnp6ennWhlacqAAXxh/E/06WYmTqXzAw/h8IRv5VdVu0QsQVQEVFRXjhhRfQ29uLc+fO4aGHHsI3v/lNDAwMAAAOHDiAU6dOoaWlBV1dXfj444/x+OOPm8+fmZlBVVUVIpEI3nnnHbz66qs4ceIEDh06tLxrRSuGwIKx6VzMnoieQQauzWhKe6KVI65DsMceeyzm8c9+9jM0NTWhu7sbRUVFeOWVV3Dy5Ek89NBDAIDjx4/j/vvvR3d3NyoqKvCnP/0JFy5cwJ///Gc4nU488MAD+MlPfoIf/vCH+PGPfwybzbZ8a0Yrg0RRaH8fwL8hinRkWQ3kZ8x9+EV3nls+BzQzM4OWlhZMTExA13X09vZiamoKHo/HrFm/fj3cbjd8Ph8qKirg8/mwYcMGOJ1Os6ayshJ79+7FwMAAvvzlL8fVw6VLl5Cdnf3FhZRQ//jHP+Ydi0oU3b5mSGYvJmYcuMs2jL9dvrjg8i5fvgyHw7HcbVISjY+PL6ou7gDq7++HruuYnJxEdnY2WltbUVJSgr6+PthsNuTm5sbUO51O+P1+AIDf748Jn9nx2bH5hMNhhMNh83EoFAIAGIbB80crwOe/NjcSAf7bNwhgcFHLEhGMj48veJmeVr6JiYlF1cUdQPfddx/6+vpgGAb+8Ic/oKamBl1dXXE3GI+GhgY8//zzN80vLy+HpvF8gmoul2vZlmWxWFBaWgpd15dtmZR8szsJXyTun4S22Wy45557UFZWhoaGBmzatAm//OUv4XK5EIlEbvrOFQgEzDeoy+W66arY7OOF3sT19fUwDMOchoeH422biFagJd+KEY1GEQ6HUVZWhoyMDHR0dJhjg4ODGBoaMr+b6bqO/v5+BIP//Nvg7e3t0DQNJSUl876G3W43L/3PTkSU+uI6BKuvr8ejjz4Kt9uNsbExnDx5Em+//TbefPNNOBwO7Nq1C3V1dcjLy4Omadi3bx90XUdFRQUAYPv27SgpKcETTzyBxsZG+P1+PPvss/B6vbDb7QlZQSJaueIKoGAwiO9973u4cuUKHA4HNm7ciDfffBOPPPIIAODFF1+E1WpFdXU1wuEwKisr8fLLL5vPT0tLQ1tbG/bu3Qtd17F69WrU1NTgyJEjy7tWlFSZmZnLtldqtVqRlpa2LMuilc8iIqK6iXiFQiE4HA4YhsHDsRVgZGRk0Vc9FqOgoIB7xClusZ9R3gtGS5aXl4e8vDzVbVAK4u8DIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMqkq27gVogIACAUCinuhIjmMvvZnP2sziclA+iTTz4BABQXFyvuhIgWMjY2BofDMe94SgZQXl4eAGBoaGjBlaNYoVAIxcXFGB4ehqZpqttJCdxmt0ZEMDY2hsLCwgXrUjKArNbPTl05HA6+KW6BpmncbnHiNovfYnYOeBKaiJRhABGRMikZQHa7HYcPH4bdblfdSkrhdosft1liWeSLrpMRESVISu4BEdHtgQFERMowgIhIGQYQESmTkgF09OhRrF27FpmZmSgvL0dPT4/qlpRpaGjAli1bkJOTg4KCAuzYsQODg4MxNZOTk/B6vcjPz0d2djaqq6sRCARiaoaGhlBVVYWsrCwUFBTg4MGDmJ6eTuaqKPPCCy/AYrFg//795jxusySRFNPc3Cw2m01+85vfyMDAgOzevVtyc3MlEAiobk2JyspKOX78uJw/f176+vrkG9/4hrjdbhkfHzdrnnrqKSkuLpaOjg45d+6cVFRUyIMPPmiOT09PS2lpqXg8HvnrX/8qp0+fljVr1kh9fb2KVUqqnp4eWbt2rWzcuFGeeeYZcz63WXKkXABt3bpVvF6v+XhmZkYKCwuloaFBYVcrRzAYFADS1dUlIiKjo6OSkZEhLS0tZs3FixcFgPh8PhEROX36tFitVvH7/WZNU1OTaJom4XA4uSuQRGNjY3LvvfdKe3u7fPWrXzUDiNsseVLqECwSiaC3txcej8ecZ7Va4fF44PP5FHa2chiGAeCfN+z29vZiamoqZputX78ebrfb3GY+nw8bNmyA0+k0ayorKxEKhTAwMJDE7pPL6/WiqqoqZtsA3GbJlFI3o169ehUzMzMxX3QAcDqduHTpkqKuVo5oNIr9+/dj27ZtKC0tBQD4/X7YbDbk5ubG1DqdTvj9frNmrm06O3Y7am5uxnvvvYd33333pjFus+RJqQCihXm9Xpw/fx5/+ctfVLeyog0PD+OZZ55Be3s7MjMzVbdzR0upQ7A1a9YgLS3tpqsRgUAALpdLUVcrQ21tLdra2vDWW2+hqKjInO9yuRCJRDA6OhpT//lt5nK55tyms2O3m97eXgSDQXzlK19Beno60tPT0dXVhZdeegnp6elwOp3cZkmSUgFks9lQVlaGjo4Oc140GkVHRwd0XVfYmToigtraWrS2tqKzsxPr1q2LGS8rK0NGRkbMNhscHMTQ0JC5zXRdR39/P4LBoFnT3t4OTdNQUlKSnBVJoocffhj9/f3o6+szp82bN2Pnzp3m/7nNkkT1WfB4NTc3i91ulxMnTsiFCxdkz549kpubG3M14k6yd+9ecTgc8vbbb8uVK1fM6dNPPzVrnnrqKXG73dLZ2Snnzp0TXddF13VzfPaS8vbt26Wvr0/eeOMNueuuu+6oS8qfvwomwm2WLCkXQCIiv/rVr8TtdovNZpOtW7dKd3e36paUATDndPz4cbPm2rVr8vTTT8uXvvQlycrKkm9961ty5cqVmOVcvnxZHn30UVm1apWsWbNGfvCDH8jU1FSS10adGwOI2yw5+Os4iEiZlDoHRES3FwYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkzP8D6QZ1mCXpm7MAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "200.0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "play(True)[-1]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
